from mcp.server.fastmcp import FastMCP
import random
from typing import Union

# 创建MCP服务实例
mcp = FastMCP("大数模素性检测服务")

def is_prime(n: Union[int, str], k: int = 10) -> bool:
    """使用Miller-Rabin算法检测大整数是否为素数
    
    Args:
        n: 待检测的整数（可以是整数或数字字符串）
        k: 检测轮数，轮数越多准确率越高
    
    Returns:
        bool: 是否为素数
    """
    # 处理字符串输入
    if isinstance(n, str):
        n = int(n)
        
    # 处理小数字的情况
    if n <= 1:
        return False
    elif n <= 3:
        return True
    elif n % 2 == 0:
        return False
    
    # 将n-1分解为d*2^s
    d = n - 1
    s = 0
    while d % 2 == 0:
        d //= 2
        s += 1
    
    # Miller-Rabin测试
    for _ in range(k):
        a = random.randint(2, min(n - 2, 1 << 20))  # 随机选择一个a
        x = pow(a, d, n)  # 计算a^d mod n
        
        if x == 1 or x == n - 1:
            continue
        
        for __ in range(s - 1):
            x = pow(x, 2, n)
            if x == n - 1:
                break
        else:
            # 如果经过s-1次平方后仍不等于n-1，则n是合数
            return False
    
    # 通过所有测试，n很可能是素数
    return True

@mcp.tool()
def check_prime(number: Union[int, str], rounds: int = 10) -> str:
    """检测一个大整数是否为素数
    
    Args:
        number: 待检测的整数（可以是整数或数字字符串）
        rounds: 检测轮数，轮数越多准确率越高，默认10轮
    
    Returns:
        str: 检测结果，包含是否为素数的判断和置信度
    """
    try:
        # 验证输入
        if isinstance(number, str):
            if not number.isdigit():
                return f"错误：输入'{number}'不是有效的整数"
            number = int(number)
            
        if number < 0:
            return "错误：不支持负数检测"
            
        # 执行素性检测
        result = is_prime(number, rounds)
        confidence = (1 - (1/4)**rounds) * 100
        
        if result:
            return f"{number} 是素数（置信度：{confidence:.6f}%）"
        else:
            return f"{number} 不是素数"
    except Exception as e:
        return f"检测失败：{str(e)}"

if __name__ == "__main__":
    print("大数模素性检测服务启动...")
    mcp.run(transport="stdio")